"""模型工厂."""

from __future__ import annotations

from typing import Dict, Type

from .base import BaseQuantModel, ModelConfig
from .dl_models import (
    GRUModelWrapper,
    LSTMModelWrapper,
    MLPModelWrapper,
    TemporalCNNModelWrapper,
)
from .ml_models import (
    CatBoostModel,
    DecisionTreeModel,
    LightGBMModel,
    LinearRegressionModel,
    LogisticRegressionModel,
    RandomForestModel,
    SVMModel,
    XGBoostModel,
)

_MODEL_REGISTRY: Dict[str, Type[BaseQuantModel]] = {
    "logistic_regression": LogisticRegressionModel,
    "linear_regression": LinearRegressionModel,
    "xgboost": XGBoostModel,
    "lightgbm": LightGBMModel,
    "random_forest": RandomForestModel,
    "svm": SVMModel,
    "decision_tree": DecisionTreeModel,
    "lstm": LSTMModelWrapper,
    "gru": GRUModelWrapper,
    "temporal_cnn": TemporalCNNModelWrapper,
    "cnn_1d": TemporalCNNModelWrapper,  # 向后兼容
    "catboost": CatBoostModel,
    "mlp": MLPModelWrapper,
}


class ModelFactory:
    """模型工厂，根据名称创建模型实例."""

    @staticmethod
    def create(model_name: str, config: ModelConfig | None = None) -> BaseQuantModel:
        """创建模型实例.

        Parameters
        ----------
        model_name:
            模型名称，支持：logistic_regression, linear_regression, xgboost,
            lightgbm, random_forest, svm, decision_tree, lstm, gru,
            temporal_cnn（或 cnn_1d），catboost, mlp
        config:
            模型配置，如果为 None 则使用默认配置

        Returns
        -------
        模型实例

        Raises
        ------
        ValueError:
            如果模型名称不支持
        """
        model_name_lower = model_name.lower()
        if model_name_lower not in _MODEL_REGISTRY:
            raise ValueError(
                f"不支持的模型名称: {model_name}。"
                f"支持的模型: {list(_MODEL_REGISTRY.keys())}"
            )

        model_class = _MODEL_REGISTRY[model_name_lower]
        if config is None:
            return model_class()
        return model_class(config)

    @staticmethod
    def register(model_name: str, model_class: Type[BaseQuantModel]) -> None:
        """注册新模型.

        Parameters
        ----------
        model_name:
            模型名称
        model_class:
            模型类
        """
        _MODEL_REGISTRY[model_name.lower()] = model_class

    @staticmethod
    def list_models() -> list[str]:
        """列出所有已注册的模型名称."""
        return list(_MODEL_REGISTRY.keys())

